import json
import time

import fastdeploy as fd
import cv2
from fastdeploy.vision.utils import mask_to_json
from yolov5.predict import main
from openvino.inference_engine import IECore

ie = IECore()
ie.set_config({'CPU_EXTENSION': ''}, 'CPU')
option = fd.RuntimeOption()
option.use_cpu()
#option.use_openvino_backend() # 一行命令切换使用 OpenVINO部署
"""
其实PaddlePaddle框架本身就可以用来做推理，只需要使用model.eval()改成评估模式就可以了。
但是实际部署时候，一般会使用Paddle Inference，因为其针对不同平台不同的应用场景进行了深度的适配优化，做到高吞吐、低时延。
总得来说，就是Paddle Inference针对推理做了很多优化。
"""

def detection_to_json(result):
    masks = []
    for mask in result.masks:
        masks.append(mask_to_json(mask))
    r_json = {
        "boxes": result.boxes,
        "scores": result.scores,
        "label_ids": result.label_ids,
        "masks": masks,
        "contain_masks": result.contain_masks
    }
    return json.dumps(r_json)


def predict(im, model='yolov5'):
    if model=='yolov5':
        result = main(im)
        return result
    elif model=='ssd':
        model = fd.vision.detection.SSD('ssd/ssd_mobilenet_v1_300_120e_voc/model.pdmodel', 'ssd/ssd_mobilenet_v1_300_120e_voc/model.pdiparams', 'ssd/ssd_mobilenet_v1_300_120e_voc/infer_cfg.yml', runtime_option=option)
    else:
        model = fd.vision.detection.FasterRCNN('fast_rcnn/faster_rcnn_r50_fpn_1x_coco/model.pdmodel', 'fast_rcnn/faster_rcnn_r50_fpn_1x_coco/model.pdiparams', 'fast_rcnn/faster_rcnn_r50_fpn_1x_coco/infer_cfg.yml', runtime_option=option)

    s = time.time()
    result = model.predict(im)
    x = detection_to_json(result)
    e = time.time()
    print("ssd_predict_time : {}s".format(e - s))
    res = json.loads(x)
    if res['boxes'] == []:
        return [im, 0]
    else:
        box = res['boxes'][0]
        p1, p2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))

        lw = max(round(sum(im.shape) / 2 * 0.003), 2)
        tf = max(lw - 1, 1)  # font thickness
        cv2.rectangle(im, p1, p2, (128, 128, 128), thickness=lw, lineType=cv2.LINE_AA)
        w, h = cv2.getTextSize('white_diploma', 0, fontScale=lw / 3, thickness=tf)[0]  # text width, height
        outside = p1[1] - h - 3 >= 0  # label fits outside box
        p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3
        cv2.rectangle(im, p1, p2, (128, 128, 128), -1, cv2.LINE_AA)  # filled
        cv2.putText(im, 'white_diploma', (p1[0], p1[1] - 2 if outside else p1[1] + h + 2), 0, lw / 3, (255, 255, 255),
                            thickness=tf, lineType=cv2.LINE_AA)
        # cv2.imwrite('./test.jpg',im)
        return [im, 1]
